package p.ithorns.framework.common.geo;

import org.locationtech.jts.algorithm.locate.IndexedPointInAreaLocator;
import org.locationtech.jts.geom.*;
import org.locationtech.jts.index.ItemVisitor;
import org.locationtech.jts.index.SpatialIndex;
import org.locationtech.jts.index.quadtree.Quadtree;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * SpatialIndexer
 * 空间索引
 *
 * @author Ht.L
 * @date 2024-01-10 17:41
 * @since 1.0.0
 */
public class SpatialIndexer {

    private final static SpatialIndex INDEXER = new Quadtree();

    private final static Map<String, Polygon> CACHE = new ConcurrentHashMap<>();


    /**
     * 获取Polygon
     *
     * @param index -索引key
     * @return Polygon
     */
    public static Polygon getPolygon(String index) {
        return CACHE.get(index);
    }

    /**
     * 添加到索引
     *
     * @param coordinateStr -位置点字符串(带格式)
     */
    public static void add(String index, String coordinateStr) {
        Polygon polygon = GeoUtil.createPolygon(coordinateStr);
        add(index, polygon);
    }

    /**
     * 添加多边形多索引
     *
     * @param index   -索引key
     * @param polygon -多边形
     */
    public static void add(String index, Polygon polygon) {
        Envelope envelope = polygon.getEnvelopeInternal();
        synchronized (CACHE) {
            // 如果已存在，则先删除
            if (CACHE.containsKey(index)) {
                // 否则会有多个items，因为是按坐标区域和图形来添加的，这里定制为不重复
                // INDEXER内部实现是NodeBase内部实现是一个List
                INDEXER.remove(envelope, index);
            }
            INDEXER.insert(envelope, index);
            CACHE.put(index, polygon);
        }
    }

    /**
     * 删除索引
     */
    public static void remove(String index) {
        synchronized (CACHE) {
            Polygon polygon = CACHE.get(index);
            if (null != polygon) {
                Envelope envelope = polygon.getEnvelopeInternal();
                INDEXER.remove(envelope, index);
            }
            CACHE.remove(index);
        }
    }

    public static void clear() {
        synchronized (CACHE) {
            Set<Map.Entry<String, Polygon>> entries = CACHE.entrySet();
            for (Map.Entry<String, Polygon> entry : entries) {
                String key = entry.getKey();
                Polygon value = entry.getValue();
                INDEXER.remove(value.getEnvelopeInternal(), key);
            }
            CACHE.clear();
        }
    }


    public static Set<String> queryIndex(double lng, double lat) {
        Coordinate coordinate = new Coordinate(lng, lat);
        return queryIndex(coordinate);
    }

    public static Set<String> queryIndex(Coordinate coordinate) {
        Envelope searchEnvelope = new Envelope(coordinate);
        LocalItemVisitor visitor = new LocalItemVisitor(coordinate);
        INDEXER.query(searchEnvelope, visitor);
        return visitor.getIndexes();
    }

    public static Set<String> queryIndex(Point point) {
        Coordinate coordinate = point.getCoordinate();
        return queryIndex(coordinate);
    }

    public static Set<Polygon> queryPolygons(double lng, double lat) {
        Coordinate coordinate = new Coordinate(lng, lat);
        return queryPolygons(coordinate);
    }


    public static Set<Polygon> queryPolygons(Coordinate coordinate) {
        Envelope searchEnvelope = new Envelope(coordinate);
        LocalItemVisitor visitor = new LocalItemVisitor(coordinate);
        INDEXER.query(searchEnvelope, visitor);
        return visitor.getPolygons();
    }


    static class LocalItemVisitor implements ItemVisitor {

        private final Coordinate coordinate;

        private final Set<Polygon> polygons = new HashSet<>();

        private final Set<String> indexes = new HashSet<>();

        /**
         * 指定需要访问的目标位置
         * {@link Location}
         */
        private final int targetLocation;


        public LocalItemVisitor(Coordinate coordinate) {
            // 默认使用Location.INTERIOR
            this(coordinate, Location.INTERIOR);
        }

        public LocalItemVisitor(Coordinate coordinate, int targetLocation) {
            this.coordinate = coordinate;
            this.targetLocation = targetLocation;
        }

        @Override
        public void visitItem(Object item) {
            Polygon polygon = CACHE.get(item);
            IndexedPointInAreaLocator locator = new IndexedPointInAreaLocator(polygon);
            int location = locator.locate(coordinate);
            if (location == this.targetLocation) {
                polygons.add(polygon);
                indexes.add((String) item);
            }
        }

        public Set<Polygon> getPolygons() {
            return polygons;
        }

        public Set<String> getIndexes() {
            return indexes;
        }

    }

}